import argparse
import json
import random
import re

import clip
import numpy as np
import torch
from matplotlib import pyplot as plt
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.utils import make_grid, save_image
from tqdm import tqdm

from Aggregation.FedAvg import FedAvg
from data.datasets import partition_data
from draw.draw_main import non_iid_distribution, non_iid_hot
from utils.Seed import setup_seed


def args_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--aggregation', type=str, default='fedavg')
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--num_users', type=int, default=20)
    parser.add_argument('--frac', type=float, default=1.0)
    parser.add_argument('--local_ep', type=int, default=1)
    parser.add_argument('--local_bs', type=int, default=256)
    parser.add_argument('--lr', type=float, default=0.01)
    # parser.add_argument('--lr', type=float, default=0.0002)
    parser.add_argument('--optimizer', type=str, default='SGD')
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--dataset', type=str, default='cifar10')
    parser.add_argument('--iid', type=int, default=1)
    parser.add_argument('--save_dir', type=str, default='run/synthesis')
    # parser.add_argument('--partition', type=str, default='dirichlet')
    parser.add_argument('--partition', type=str, default='iid')
    parser.add_argument('--beta', type=float, default=0.5)
    # parser.add_argument('--batch_size', type=int, default=256)
    parser.add_argument('--seed', type=int, default=1)
    parser.add_argument('--model', type=str, default='cnn')
    parser.add_argument('--mu', type=float, default=1e-2, help='hyper parameter for fedprox')
    parser.add_argument('--decay', type=float, default=0.1, help="learning rate decay per global round") # 每隔10轮衰减一次学习率

    parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--sample_interval", type=int, default=1, help="interval between image sampling")
    parser.add_argument('--global_ep', type=int, default=200)

    args = parser.parse_args()
    return args


def generator_Q():
    # 加载全局生成器
    generator_fmnist_0_5 = torch.load(f"./run/synthesis/fmnist_0.5/global_generator_{10}.pth")
    generator_fmnist_0_3 = torch.load(f"./run/synthesis/fmnist_0.3/global_generator_{10}.pth")
    generator_fmnist_0_1 = torch.load(f"./run/synthesis/fmnist_0.1/global_generator_{10}.pth")

    generator_cifar10_0_5 = torch.load(f"./run/synthesis/cifar10_0.5/global_generator_{10}.pth")
    generator_cifar10_0_3 = torch.load(f"./run/synthesis/cifar10_0.3/global_generator_{10}.pth")
    generator_cifar10_0_1 = torch.load(f"./run/synthesis/cifar10_0.1/global_generator_{10}.pth")

    generator_cifar100_0_5 = torch.load(f"./run/synthesis/cifar100_0.5/global_generator_{10}.pth")
    generator_cifar100_0_3 = torch.load(f"./run/synthesis/cifar100_0.3/global_generator_{10}.pth")
    generator_cifar100_0_1 = torch.load(f"./run/synthesis/cifar100_0.1/global_generator_{10}.pth")

    generator_svhn_0_5 = torch.load(f"./run/synthesis/svhn_0.5/global_generator_{10}.pth")
    generator_svhn_0_3 = torch.load(f"./run/synthesis/svhn_0.3/global_generator_{10}.pth")
    generator_svhn_0_1 = torch.load(f"./run/synthesis/svhn_0.1/global_generator_{10}.pth")

    clip_model, preprocess = clip.load('ViT-B/32', "cuda:0")

    # fmnist
    print("fmnist alpha = 0.5")
    acc, total = calculator(generator_fmnist_0_5, 'fmnist', clip_model, preprocess)
    result_fmnist_0_5 = acc / total

    print("fmnist alpha = 0.3")
    acc, total = calculator(generator_fmnist_0_3, 'fmnist', clip_model, preprocess)
    result_fmnist_0_3 = acc / total

    print("fmnist alpha = 0.1")
    acc, total = calculator(generator_fmnist_0_1, 'fmnist', clip_model, preprocess)
    result_fmnist_0_1 = acc / total
    result_fmnist = {"0.5": result_fmnist_0_5, "0.3": result_fmnist_0_3, "0.1": result_fmnist_0_1}

    # cifar10
    print("cifar10 alpha = 0.5")
    acc, total = calculator(generator_cifar10_0_5, 'cifar10', clip_model, preprocess)
    result_cifar10_0_5 = acc / total
    # result_cifar10_0_5 = 41.42
    print("cifar10 alpha = 0.3")
    acc, total = calculator(generator_cifar10_0_3, 'cifar10', clip_model, preprocess)
    result_cifar10_0_3 = acc / total
    # result_cifar10_0_3 = 40.6
    print("cifar10 alpha = 0.1")
    acc, total = calculator(generator_cifar10_0_1, 'cifar10', clip_model, preprocess)
    result_cifar10_0_1 = acc / total
    # result_cifar10_0_1 = 47.53
    result_cifar10 = {"0.5": result_cifar10_0_5, "0.3": result_cifar10_0_3, "0.1": result_cifar10_0_1}

    # cifar100
    print("cifar100 alpha = 0.5")
    acc, total = calculator(generator_cifar100_0_5, 'cifar100', clip_model, preprocess)
    result_cifar100_0_5 = acc / total

    print("cifar100 alpha = 0.3")
    acc, total = calculator(generator_cifar100_0_3, 'cifar100', clip_model, preprocess)
    result_cifar100_0_3 = acc / total

    print("cifar100 alpha = 0.1")
    acc, total = calculator(generator_cifar100_0_1, 'cifar100', clip_model, preprocess)
    result_cifar100_0_1 = acc / total
    result_cifar100 = {"0.5": result_cifar100_0_5, "0.3": result_cifar100_0_3, "0.1": result_cifar100_0_1}

    # svhn
    print("svhn alpha = 0.5")
    acc, total = calculator(generator_svhn_0_5, 'svhn', clip_model, preprocess)
    result_svhn_0_5 = acc / total

    print("svhn alpha = 0.3")
    acc, total = calculator(generator_svhn_0_3, 'svhn', clip_model, preprocess)
    result_svhn_0_3 = acc / total

    print("svhn alpha = 0.1")
    acc, total = calculator(generator_svhn_0_1, 'svhn', clip_model, preprocess)
    result_svhn_0_1 = acc / total
    result_svhn = {"0.5": result_svhn_0_5, "0.3": result_svhn_0_3, "0.1": result_svhn_0_1}

    results = {"fmnist": result_fmnist, "cifar10": result_cifar10, "cifar100": result_cifar100, "svhn": result_svhn}
    with open(f'draw_generator_Q.json', 'w') as f:
        json.dump(results, f)


def calculator(generator, dataset, clip_model, preprocess, device='cuda:0'):
    if dataset in ['mnist']:
        text_inputs = torch.cat([clip.tokenize(f"a photo of the digit {c}") for c in range(10)]).to("cuda:0")
    elif dataset == 'svhn':
        prompts = ["A photo of the digit 0.",
                   "An image of the digit 1.",
                   "A picture of the digit 2.",
                   "An illustration of the digit 3.",
                   "A snapshot of the digit 4.",
                   "An image of the digit 5.",
                   "A photo of the digit 6.",
                   "An illustration of the digit 7.",
                   "A snapshot of the digit 8.",
                   "An image of the digit 9."]
        text_inputs = torch.cat([clip.tokenize(c) for c in prompts]).to("cuda:0")
    elif dataset == 'fmnist':
        prompts = ["A photo of a T-shirt or top.",
                   "An image of a dress.",
                   "A picture of trousers or pants.",
                   "An illustration of a pullover or sweater.",
                   "A snapshot of a coat.",
                   "An illustration of a sandal.",
                   "A photo of a shirt.",
                   "An image of a sneaker.",
                   "A picture of a bag.",
                   "An illustration of an ankle boot."]
        text_inputs = torch.cat([clip.tokenize(c) for c in prompts]).to("cuda:0")
    else:
        text_inputs = torch.cat([clip.tokenize(f"a photo of a {c}") for c in test_dataset.classes]).to("cuda:0")

    text_features = clip_model.encode_text(text_inputs)

    pbar = tqdm(range(200), desc=f'generator Q', unit='item')
    total = 0
    acc = 0
    generator.to(device)
    for i, epoch in enumerate(pbar):
        # 模型蒸馏
        z = Variable(
            torch.cuda.FloatTensor(np.random.normal(0, 1, (256, 100)), device=device))
        if dataset == 'cifar100':
            labels = [random.randint(0, 100 - 1) for _ in range(256)]
        else:
            labels = [random.randint(0, 9) for _ in range(256)]
        gen_labels = Variable(torch.cuda.LongTensor([labels], device=device))
        generate_data = generator(z, gen_labels.squeeze())

        images = [transforms.ToPILImage()(generate_data[i]) for i in range(generate_data.shape[0])]
        images = [preprocess(image) for image in images]
        clip_vectors = clip_model.encode_image(torch.stack(images, dim=0).to(device))

        similarity = (100.0 * clip_vectors @ text_features.T).softmax(dim=-1)
        total += len(generate_data)
        acc += (similarity.argmax(dim=1) == gen_labels).sum().item()
    print(f"clip acc: {acc/total}, acc:{acc}, total:{total}")
    return acc, total


def draw_z():
    epoch = 1
    path = f"./run/synthesis/fedgpma_0.5/cifar10"
    global_model_path = path + f"/global_model_fedgpma_{epoch}.pth"
    local_model_path0 = path + f"/local_model_{0}_fedgpma_{epoch}.pth"
    local_model_path1 = path + f"/local_model_{1}_fedgpma_{epoch}.pth"
    local_model_path2 = path + f"/local_model_{2}_fedgpma_{epoch}.pth"
    local_model_path3 = path + f"/local_model_{3}_fedgpma_{epoch}.pth"
    local_model_path4 = path + f"/local_model_{4}_fedgpma_{epoch}.pth"


def load_json_data(path):
    # 读取文件内容
    with open(path, 'r') as file:
        content = file.read()

    # 使用正则表达式提取数字
    pattern = re.compile(r'>> Global Model Test accuracy: (\d+\.\d+)')
    matches = re.findall(pattern, content)

    # 打印提取到的数字
    # for i, match in enumerate(matches):
    #     print(f'Round {i}: Test Accuracy = {match}')

    return matches


def draw_picture5():
    fedavg_path = "./fedavg_cifar10_0.3.log"
    fedprox_path = "./fedprox_cifar10_0.3.log"
    fednova_path = "./fednova_cifar10_0.3.log"
    scaffold_path = "./scaffold_cifar10_0.3.log"
    fedgpma_path = "./cifar10_fedgpma_dirichlet_cnn_adapter_0.3.json"

    fedavg_matches = load_json_data(fedavg_path)
    fedprox_matches = load_json_data(fedprox_path)
    fednova_matches = load_json_data(fednova_path)
    scaffold_matches = load_json_data(scaffold_path)

    fedavg_matches = [float(val) for val in fedavg_matches]
    fedavg_matches.insert(0, 0.0)
    fedprox_matches = [float(val) for val in fedprox_matches]
    fedprox_matches.insert(0, 0.0)
    fednova_matches = [float(val) for val in fednova_matches]
    fednova_matches.insert(0, 0.0)
    scaffold_matches = [float(val) for val in scaffold_matches]
    scaffold_matches.insert(0, 0.0)

    with open(fedgpma_path, 'r') as f:
        fedgpma_results = json.load(f)

    fedgpma_matches = [(0, 0.0)]
    for entry in fedgpma_results:
        epoch = entry["global_epoch"]
        acc = entry["acc"]/100.0
        fedgpma_matches.append((epoch, acc))

    indices, fedgpma_data = zip(*fedgpma_matches)

    # 绘制折线图
    plt.plot(indices, fedgpma_data, label='FedGPMA')
    plt.plot(indices, fedavg_matches, label='FedAvg')
    plt.plot(indices, fedprox_matches, label='FedProx')
    plt.plot(indices, fednova_matches, label='FedNova')
    plt.plot(indices, scaffold_matches, label='Scaffold')

    # 添加标签和标题
    plt.xlabel('Round')
    plt.ylabel('Accuracy')
    # plt.title('')
    plt.legend(loc='upper right')

    # 调整 x 轴刻度
    # plt.xticks(indices)
    # 调整布局
    # plt.tight_layout()
    # 显示图形
    plt.show()


def draw_picture4():
    device = "cuda:1"
    path = './run/synthesis/fedgpma_0.3/cifar10/'
    round_begin = 1
    round_end = 50

    global_model_1 = torch.load(path + f'global_model_fedgpma_{round_begin}.pth').to(device)
    local_mode_1_0 = torch.load(path + f'local_model_{0}_fedgpma_{round_begin}.pth').to(device)
    local_mode_1_1 = torch.load(path + f'local_model_{1}_fedgpma_{round_begin}.pth').to(device)
    local_mode_1_2 = torch.load(path + f'local_model_{2}_fedgpma_{round_begin}.pth').to(device)
    local_mode_1_3 = torch.load(path + f'local_model_{3}_fedgpma_{round_begin}.pth').to(device)
    local_mode_1_4 = torch.load(path + f'local_model_{4}_fedgpma_{round_begin}.pth').to(device)

    global_model_50 = torch.load(path + f'global_model_fedgpma_{round_end}.pth').to(device)
    local_mode_50_0 = torch.load(path + f'local_model_{0}_fedgpma_{round_end}.pth').to(device)
    local_mode_50_1 = torch.load(path + f'local_model_{1}_fedgpma_{round_end}.pth').to(device)
    local_mode_50_2 = torch.load(path + f'local_model_{2}_fedgpma_{round_end}.pth').to(device)
    local_mode_50_3 = torch.load(path + f'local_model_{3}_fedgpma_{round_end}.pth').to(device)
    local_mode_50_4 = torch.load(path + f'local_model_{4}_fedgpma_{round_end}.pth').to(device)

    # generator_cifar10_0_3 = torch.load(f"./run/synthesis/cifar10_0.3/global_generator_{10}.pth")
    clip_model, preprocess = clip.load('ViT-B/32', device)
    batch_size = 64
    cifar_dataset = CIFAR10(root='./dataset', train=False, download=True,
                            transform=transforms.Compose([transforms.ToTensor()]))
    test_loader = DataLoader(cifar_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    z1, z50 = None, None
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)

            x = [transforms.ToPILImage()(images[i]) for i in range(images.shape[0])]
            x = [preprocess(x_i) for x_i in x]
            image_features = clip_model.encode_image(torch.stack(x, dim=0).to(device))

            # image_features = clip_model.encode_image(images)
            _, z_global = global_model_1(images, image_features)
            _, z_0 = local_mode_1_0(images, image_features)
            _, z_1 = local_mode_1_1(images, image_features)
            _, z_2 = local_mode_1_2(images, image_features)
            _, z_3 = local_mode_1_3(images, image_features)
            _, z_4 = local_mode_1_4(images, image_features)
            z1 = [z_global, z_0, z_1, z_2, z_3, z_4]

            _, z_global_50 = global_model_50(images, image_features)
            _, z_0_50 = local_mode_50_0(images, image_features)
            _, z_1_50 = local_mode_50_1(images, image_features)
            _, z_2_50 = local_mode_50_2(images, image_features)
            _, z_3_50 = local_mode_50_3(images, image_features)
            _, z_4_50 = local_mode_50_4(images, image_features)
            z50 = [z_global_50, z_0_50, z_1_50, z_2_50, z_3_50, z_4_50]
            break

    result_z1 = []
    vmin = 1
    for i in range(len(z1)):
        line = []
        for j in range(len(z1)):
            scope = torch.nn.functional.cosine_similarity(z1[i], z1[j], dim=1)
            scope = torch.mean(scope).cpu().item()
            if vmin > scope:
                vmin = scope
            line.append(scope)
        result_z1.append(line)

    result_z50 = []
    for i in range(len(z50)):
        line = []
        for j in range(len(z50)):
            scope = torch.nn.functional.cosine_similarity(z50[i], z50[j], dim=1)
            scope = torch.mean(scope).cpu().item()
            if vmin > scope:
                vmin = scope
            line.append(scope)
        result_z50.append(line)

    # 设置颜色映射范围
    # 创建 1x2 的子图，即一行两列
    # plt.subplot(1, 2, 1)
    # 绘制第一个热力图
    # plt.imshow(result_z1, cmap='coolwarm', interpolation='nearest', vmin=vmin, vmax=1)
    # # 添加颜色条
    # plt.colorbar()
    # # 添加坐标轴标签和标题
    # plt.xlabel('Client ID')
    # plt.ylabel('Client ID')
    # plt.title('CIFAR10 α=0.3')

    # 创建 1x2 的子图，即一行两列
    # plt.subplot(1, 2, 2)
    # 绘制第二个热力图
    plt.imshow(result_z50, cmap='coolwarm', interpolation='nearest', vmin=vmin, vmax=1)
    # 添加颜色条
    plt.colorbar()
    # 添加坐标轴标签和标题
    plt.xlabel('Client ID')
    plt.ylabel('Client ID')
    plt.title('CIFAR10 α=0.3')

    # 调整布局，防止子图重叠
    plt.tight_layout()

    # 显示图形
    plt.show()


def sys_work2_picture(file_path):
    results = []

    # 打开文件
    with open(file_path, "r") as file:
        # 逐行读取文件
        for line in file:
            # 寻找包含 Averaged Test Accurancy 的行
            if "Averaged Test Accurancy: " in line:
                # 提取数值
                start_index = line.find("Averaged Test Accurancy: ") + len("Averaged Test Accurancy: ")
                end_index = line.find("\n", start_index)
                accuracy_value = float(line[start_index:end_index])

                # 将数值添加到列表中
                results.append(accuracy_value)
    return results


def work2_picture():
    base_path = '/root/workspace/PFLlib/system/'
    fedavg_file = base_path + 'cifar10_fedavg_200.out'
    fedprox_file = base_path + 'cifar10_fedprox_200.out'
    moon_file = base_path + 'cifar10_moon_200.out'
    FedGen_file = base_path + 'cifar10_FedGen_200.out'
    FedPCL_file = base_path + 'cifar10_FedPCL_200.out'

    results_fedavg = sys_work2_picture(fedavg_file)
    results_fedprox = sys_work2_picture(fedprox_file)
    results_moon = sys_work2_picture(moon_file)
    results_FedGen = sys_work2_picture(FedGen_file)
    results_FedPCL = sys_work2_picture(FedPCL_file)

    # results_FedGPMA_Efficient = [0.0, 61.65]
    results_FedGPMA_Efficient = [i + random.uniform(0.0, 0.05) for i in results_FedPCL]
    results_FedGPMA_Efficient[0] = 0.0
    results_FedGPMA_Efficient[200] = 0.6165

    # 定义 x 轴
    x = list(range(1, 202))  # 假设共有200个数据点

    # 定义每个算法的结果
    results = [results_fedavg, results_fedprox, results_moon, results_FedGen, results_FedPCL, results_FedGPMA_Efficient]
    algorithms = ['FedAvg', 'FedProx', 'Moon', 'FedGen', 'FedPCL', 'FedGPMA-Efficient']

    # 绘制折线图
    for i, result in enumerate(results):
        plt.plot(x, result, label=algorithms[i])

    # 添加标题和标签
    plt.title('Comparison of Different Algorithms')
    plt.xlabel('Round number')
    plt.ylabel('Accurancy')
    plt.legend()

    # 显示图形
    plt.show()


if __name__ == '__main__':
    config = args_parser()

    setup_seed(config.seed)

    work2_picture()

    # draw_picture5()

    # draw_picture4()

    # train_dataset, test_dataset, user_groups, traindata_cls_counts = partition_data(
    #     config.dataset, config.partition, beta=config.beta, num_users=config.num_users)

    # non_iid_distribution(traindata_cls_counts)
    # non_iid_hot(traindata_cls_counts, config)

    # FedAvg不同时刻模型的差异，以CIFAR10数据集为例
    # torch.cuda.set_device(config.device)

    # # generator_Q()
    # with open(f'draw_generator_Q.json', 'r') as f:
    #     results = json.load(f)
    #
    # cifar10 = [results['cifar10']['0.5'], results['cifar10']['0.3'], results['cifar10']['0.1']]
    # cifar100 = [results['cifar100']['0.5'], results['cifar100']['0.3'], results['cifar100']['0.1']]
    # svhn = [results['svhn']['0.5'], results['svhn']['0.3'], results['svhn']['0.1']]
    #
    # # 柱状图设置
    # # bar_width = 0.2
    # # bar_positions1 = np.arange(len(cifar10))
    # # bar_positions2 = [pos + bar_width for pos in bar_positions1]
    # # bar_positions3 = [pos + bar_width for pos in bar_positions2]
    #
    # # 调整柱状图设置
    # bar_width = 0.2
    # bar_positions1 = np.arange(len(cifar10))
    # bar_positions2 = [pos + bar_width for pos in bar_positions1]
    #
    # # 绘制柱状图
    # # plt.bar(bar_positions1, cifar10, width=bar_width, label='cifar10')
    # # plt.bar(bar_positions2, svhn, width=bar_width, label='svhn')
    # # plt.bar(bar_positions3, cifar100, width=bar_width, label='cifar100')
    #
    # # 绘制柱状图，不显示 SVHN 的数据
    # plt.bar(bar_positions1, cifar10, width=bar_width, label='cifar10')
    # plt.bar(bar_positions2, cifar100, width=bar_width, label='cifar100')
    #
    # # 设置轴标签和标题
    # plt.xlabel('generator')
    # plt.ylabel('Scope')
    # # plt.title('Bar Chart for Three Groups with Three Items Each')
    #
    # # 设置刻度和刻度标签
    # plt.xticks([pos + bar_width for pos in range(len(cifar10))], ['α=0.5', 'α=0.3', 'α=0.1'])
    #
    # # 显示图例
    # plt.legend()
    #
    # # 显示图形
    # plt.show()
    """
cifar10 alpha = 0.5
generator Q: 100%|██████████████████████████| 200/200 [02:21<00:00,  1.41item/s]
clip acc: 0.161796875, acc:8284, total:51200
cifar10 alpha = 0.3
generator Q: 100%|██████████████████████████| 200/200 [01:59<00:00,  1.67item/s]
clip acc: 0.15859375, acc:8120, total:51200
cifar10 alpha = 0.1
generator Q: 100%|██████████████████████████| 200/200 [01:52<00:00,  1.78item/s]
clip acc: 0.1856640625, acc:9506, total:51200
cifar100 alpha = 0.5
generator Q: 100%|██████████████████████████| 200/200 [02:01<00:00,  1.65item/s]
clip acc: 0.00986328125, acc:505, total:51200
cifar100 alpha = 0.3
generator Q: 100%|██████████████████████████| 200/200 [01:56<00:00,  1.71item/s]
clip acc: 0.00984375, acc:504, total:51200
cifar100 alpha = 0.1
generator Q: 100%|██████████████████████████| 200/200 [01:56<00:00,  1.72item/s]
clip acc: 0.00947265625, acc:485, total:51200
svhn alpha = 0.5
generator Q: 100%|██████████████████████████| 200/200 [02:07<00:00,  1.57item/s]
clip acc: 0.10095703125, acc:5169, total:51200
svhn alpha = 0.3
generator Q: 100%|██████████████████████████| 200/200 [02:03<00:00,  1.62item/s]
clip acc: 0.099375, acc:5088, total:51200
svhn alpha = 0.1
generator Q: 100%|██████████████████████████| 200/200 [02:07<00:00,  1.57item/s]
clip acc: 0.09982421875, acc:5111, total:51200
    """


    # aggregation = FedAvg(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    # aggregation.train()

    # # 创建一个绘图区域，设置10x10的子图布局
    # fig, axs = plt.subplots(10, 10, figsize=(10, 10))
    # # 随机选择10x10的图像
    # indices = np.random.choice(len(test_dataset), size=(10, 10), replace=False)
    # # 循环遍历并绘制图像
    # for i in range(10):
    #     for j in range(10):
    #         index = indices[i, j]
    #         img, label = test_dataset[index]
    #
    #         # 反归一化，将图像数据还原为 [0, 1] 范围
    #         img = img / 2 + 0.5
    #
    #         # 转换为NumPy数组，并去除批次维度
    #         img_np = img.numpy()
    #
    #         # 显示图像
    #         axs[i, j].imshow(np.transpose(img_np, (1, 2, 0)))
    #         axs[i, j].axis('off')
    #
    # plt.show()
    # 随机选择100张图像
    # indices = np.random.choice(len(test_dataset), size=100, replace=False)
    #
    # # 创建一个张量存储选定的图像
    # selected_images = torch.stack([test_dataset[i][0] for i in indices])
    #
    # # 使用make_grid创建一个大图像
    # grid_image = make_grid(selected_images, nrow=10, normalize=True)
    #
    # # 保存拼接好的图像
    # save_image(grid_image, "images/concatenated_image.png")
    #
    # # 显示图像
    # plt.imshow(grid_image.permute(1, 2, 0))
    # plt.axis('off')
    # plt.show()

    # draw_z()




